import csv
import os
import numpy as np

'''
数据保存为CSV
'''
def write_d(data_all, file_name):
    with open(file_name+".csv", "w", newline="") as f:
        f_csv = csv.writer(f)
        # f_csv.writerow(headers)
        f_csv.writerows(data_all)


def check_path_exist(directory_path):
    if not os.path.exists(directory_path):
        os.mkdir(directory_path)

    
'''
保存所有数据
'''
def write_all(all_list, model,output_dir,file_name):
    OUTPUT_DIR = output_dir + model+"/"+file_name

    check_path_exist(output_dir)
    check_path_exist(output_dir + model)
    check_path_exist(output_dir + model+"/"+file_name)

    # 保存力
    force_name = ["force_aero","force_body","force_wind","force_rslt"]
    for i, file in enumerate(all_list[0]):
        write_d(file, OUTPUT_DIR +"/"+ force_name[i])
        
    # 保存系数
    coeff_name = ["coeff_body","coeff_wind","coeff_rslt"]
    for i, file in enumerate(all_list[1]):
        write_d(file, OUTPUT_DIR +"/"+ coeff_name[i])

    # 保存角度
    angle_name = ["angle_body","angle_wind"]
    for i, file in enumerate(all_list[2]):
        write_d(file, OUTPUT_DIR +"/"+ angle_name[i])

    print("------")
    print("保存成功")



'''
保存为数据集
Cx, Cy, Cz, Cl, Cm, Cn
α, β, p, q, r
'''
def write_dataset(all_list, model,output_dir,file_name):
    OUTPUT_DIR = output_dir + model

    check_path_exist(output_dir)
    check_path_exist(output_dir + model)
    [C_n_Q_body, Angle_body, Angle_wind] = all_list
    
    # Cx, Cy, Cz, Cl, Cm, Cn
    C_n_Q_body = C_n_Q_body[:, :6]
    
    # p, q, r
    Angle_body = Angle_body[:, [7,6,8]]
    
    combined_data = np.hstack((C_n_Q_body, Angle_wind, Angle_body))

    
    columns = ['Cx', 'Cy', 'Cz', 'Cl', 'Cm', 'Cn', 'alpha', 'beta', 'p', 'q', 'r']
    
    with open(OUTPUT_DIR +"/"+file_name+".csv", 'w', newline='') as file:
        writer = csv.writer(file)
        # Write the header
        writer.writerow(columns)
        # Write the data
        for row in combined_data:
            writer.writerow(row)

    print("------")
    print("保存成功")
